package org.jing.ext.media.m3u8;

import org.jing.core.lang.Const;
import org.jing.core.lang.EnumFactory;
import org.jing.core.lang.ExceptionHandler;
import org.jing.core.lang.JingException;
import org.jing.core.lang.data.FileContent;
import org.jing.core.lang.data.Wrapper;
import org.jing.core.logger.JingLogger;
import org.jing.core.util.ClassUtil;
import org.jing.core.util.FileUtil;
import org.jing.core.util.GenericUtil;
import org.jing.core.util.StringUtil;
import org.jing.ext.media.m3u8.tag.M3U8AllowCache;
import org.jing.ext.media.m3u8.tag.M3U8Begin;
import org.jing.ext.media.m3u8.tag.M3U8End;
import org.jing.ext.media.m3u8.tag.M3U8Key;
import org.jing.ext.media.m3u8.tag.M3U8MediaPartition;
import org.jing.ext.media.m3u8.tag.M3U8MediaSequence;
import org.jing.ext.media.m3u8.tag.M3U8PlaylistType;
import org.jing.ext.media.m3u8.tag.M3U8TargetDuration;
import org.jing.ext.media.m3u8.tag.M3U8Version;
import org.jing.ext.security.AES;
import org.jing.ext.security.Common;
import org.jing.ext.web.RequestEntity;
import org.jing.ext.web.ResponseEntity;
import org.jing.ext.web.Sender;
import org.jing.ext.web.WebConst;
import org.jing.ext.web.WebHeaders;
import org.jing.ext.web.WebUtil;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.math.BigDecimal;
import java.net.InetSocketAddress;
import java.net.Proxy;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

/**
 * Description: <br>
 *
 * @author bksmate <br>
 * @since 2024-12-24 <br>
 */
public class M3U8Factory {
    private static final JingLogger LOGGER = JingLogger.getLogger(M3U8Factory.class);

    private static final Pattern PATTERN_TAG = Pattern.compile("^#([^:]+)(?::(.*))?$");

    private static M3U8Config config = new M3U8Config();

    public static synchronized M3U8Config getRootConfig() {
        return config;
    }

    public static synchronized void setRootConfig(M3U8Config config) {
        M3U8Factory.config = config;
    }

    public static M3U8 downloadM3U8(String m3u8Url) {
        return downloadM3U8(getRootConfig(), m3u8Url);
    }

    public static M3U8 downloadM3U8(M3U8Config config, String m3u8Url) {
        // 校验输出路径
        config.setOutputPath(FileUtil.buildPath(config.getOutputPath()));
        FileUtil.mkdirs(config.getOutputPath());
        // 校验文件名
        String fileName = config.getFileName();
        // 文件名必填
        ExceptionHandler.publishIfMatch(StringUtil.isEmpty(fileName), "fileName required");
        // 不能指定文件后缀, 因为都是固定的后缀
        ExceptionHandler.publishIfMatch(StringUtil.isNotEmpty(FileUtil.getFileSuffix(fileName)), "fileName suffix should not be set");
        File file = new File(config.getOutputPath() + (StringUtil.isEmpty(config.getOutputPath()) ? "" : File.separator) + fileName + ".m3u8");
        ExceptionHandler.publishIfMatch(file.exists() && file.isFile() && !FileUtil.deleteFile(file), "failed to clear output file: [{}]", file);
        // 校验相对路径
        if (StringUtil.isEmpty(config.getRelativeUrl())) {
            config.setRelativeUrl(m3u8Url.substring(0, m3u8Url.lastIndexOf("/") + 1));
        }

        // 请求获得m3u8文件
        byte[] buffer = request(config, m3u8Url);
        String m3u8Content = new String(buffer).replaceAll("\r\n", "\n");

        // 解析
        List<String> rows = Arrays.stream(m3u8Content.split("\n")).collect(Collectors.toList());
        List<BaseM3U8Tag> tags = analysisRows(config, rows);
        M3U8 m3u8 = analysisTags(config, tags);
        m3u8.setOutputFile(file);
        LOGGER.debug(m3u8);

        // 处理加密规则
        analysisKey(config, m3u8);

        // 处理切片路径
        analysisMediaPartition(config, m3u8);

        if (config.isAsync()) {
            new AsyncDownloader(config, m3u8).start();
        }
        else {
            execute(config, m3u8);
        }

        return m3u8;
    }

    private static List<BaseM3U8Tag> analysisRows(M3U8Config config, List<String> rows) {
        LOGGER.info("analysis m3u8 content...");
        List<BaseM3U8Tag> tags = new ArrayList<>();
        boolean isHeader = true;
        for (int i$ = 0, l$ = GenericUtil.count(rows); i$ < l$; ++i$) {
            String row = rows.get(i$);
            Matcher matcher = PATTERN_TAG.matcher(row);
            // 是标签
            if (matcher.find()) {
                String tagName = matcher.group(1);
                String tagValue = matcher.groupCount() > 1 ? matcher.group(2) : "";
                M3U8TagType tagType = EnumFactory.mapEnum(M3U8TagType.class, tagName);
                if (null == tagType) {
                    LOGGER.debug("unknown tag: [{}]", tagName);
                    ExceptionHandler.publishIfMatch(!config.isSkipUnknownTags(), "unknown tag: [{}]", tagName);
                    continue;
                }
                BaseM3U8Tag tag = ClassUtil.createInstance(tagType.getTagClass());
                tag.setParameter(tagValue);
                i$ ++;
                List<String> parameters = null;
                for (; i$ < l$; i$ ++) {
                    row = rows.get(i$);
                    if (row.matches(PATTERN_TAG.pattern())) {
                        break;
                    }
                    if (null == parameters) {
                        parameters = new ArrayList<>();
                    }
                    parameters.add(row);
                }
                i$ --;
                tag.setExtParameters(parameters);
                tags.add(tag);
            }
            else if (0 == i$) {
                throw new JingException(row);
            }
        }
        return tags;
    }

    private static M3U8 analysisTags(M3U8Config config, List<BaseM3U8Tag> tags) {
        LOGGER.info("analysis m3u8 tags...");
        M3U8 m3u8 = new M3U8();
        for (int i$ = 0, l$ = tags.size(); i$ < l$; ++ i$) {
            BaseM3U8Tag tag = tags.get(i$);
            ExceptionHandler.publishIfMatch(0 == i$ && !(tag instanceof M3U8Begin), "invalid m3u8, begin not found");
            if (tag instanceof M3U8End) {
                break;
            }
            else if (tag instanceof M3U8Version) {
                m3u8.setVersion((M3U8Version) tag);
            }
            else if (tag instanceof M3U8AllowCache) {
                m3u8.setAllowCache((M3U8AllowCache) tag);
            }
            else if (tag instanceof M3U8TargetDuration) {
                m3u8.setTargetDuration((M3U8TargetDuration) tag);
            }
            else if (tag instanceof M3U8MediaSequence) {
                m3u8.setMediaSequence((M3U8MediaSequence) tag);
            }
            else if (tag instanceof M3U8PlaylistType) {
                m3u8.setPlaylistType((M3U8PlaylistType) tag);
            }
            else if (tag instanceof M3U8Key) {
                m3u8.setKey((M3U8Key) tag);
            }
            else if (tag instanceof M3U8MediaPartition) {
                m3u8.getMediaPartitions().add((M3U8MediaPartition) tag);
            }
        }
        return m3u8;
    }

    private static void analysisKey(M3U8Config config, M3U8 m3u8) {
        LOGGER.info("analysis decrypt key...");
        M3U8Key key = m3u8.getKey();
        if (null != key) {
            // key
            String keyPath = m3u8.getKey().getKeyPath();
            if (!WebUtil.isUrl(keyPath)) {
                keyPath = mergeUrl(config.getRelativeUrl(), keyPath);
            }
            byte[] keyBuffer = request(config, keyPath);
            key.setKey(keyBuffer);

            // iv
            String ivStr = key.getIvStr();
            if (ivStr.startsWith("0x")) {
                ivStr = ivStr.substring(2);
            }
            byte[] ivBuffer = Common.parseHexStr2Byte(ivStr);
            ivBuffer = null == ivBuffer || ivBuffer.length != 16 ?  new byte[16] : ivBuffer;
            key.setIv(ivBuffer);
        }
    }

    private static void analysisMediaPartition(M3U8Config config, M3U8 m3u8) {
        LOGGER.info("analysis media partitions...");
        for (M3U8MediaPartition p$ : m3u8.getMediaPartitions()) {
            if (!WebUtil.isUrl(p$.getMedia())) {
                p$.setMedia(mergeUrl(config.getRelativeUrl(), p$.getMedia()));
            }
        }
    }

    private static void execute(M3U8Config config, M3U8 m3u8) {
        if (null != config.getStartListener()) {
            config.getStartListener().execute(config, m3u8);
        }

        // 尸块下载
        File[] mediaFiles = downloadMediaPartition(config, m3u8);

        // 合并片段
        combineMediaFiles(config, m3u8, mediaFiles);

        if (null != config.getEndListener()) {
            config.getEndListener().execute(config, m3u8);
        }
    }

    private static File[] downloadMediaPartition(M3U8Config config, M3U8 m3u8) {
        List<M3U8MediaPartition> partitions = m3u8.getMediaPartitions();
        int size = GenericUtil.count(partitions);
        ExceptionHandler.publishIfMatch(0 == size, "no media partitions found");
        File[] mediaFiles = new File[size];
        int threads = Math.min(config.getThreads(), size);
        ExecutorService executorService = Executors.newFixedThreadPool(threads);
        List<Future<Void>> futures = new ArrayList<>();
        for (int i$ = 0; i$ < size; ++ i$) {
            futures.add(executorService.submit(new M3U8MediaPartitionDownloader(config, m3u8, partitions, mediaFiles, i$)));
        }
        for (Future<Void> future : futures) {
            try {
                m3u8.addCompletedCount();
                future.get();
            }
            catch (JingException e) {
                throw e;
            }
            catch (Throwable t) {
                throw new JingException(t);
            }
        }
        executorService.shutdown();
        // 以防不知道万一什么的万一做个检查
        for (int i$ = 0; i$ < size; ++ i$) {
            ExceptionHandler.publishIfMatch(null == mediaFiles[i$], "unexpected failed in {}/{}", (i$ + 1), size);
        }
        return mediaFiles;
    }

    private static void combineMediaFiles(M3U8Config config, M3U8 m3u8, File[] mediaFiles) {
        if (config.isAutoCombine()) {
            File tsFile = new File(config.getOutputPath() + (StringUtil.isEmpty(config.getOutputPath()) ? "" : File.separator) + config.getFileName() + ".ts");
            LOGGER.info("combine partition files to [{}]", tsFile.getAbsolutePath());
            System.gc();
            FileUtil.deleteFile(tsFile);
            try (FileOutputStream fos = new FileOutputStream(tsFile)) {
                byte[] buffer = new byte[4096];
                for (File f$ : mediaFiles) {
                    try (FileInputStream fis = new FileInputStream(f$)) {
                        int len;
                        while ((len = fis.read(buffer)) != -1) {
                            fos.write(buffer, 0, len);
                        }
                    }
                    FileUtil.deleteFile(f$);
                }
            }
            catch (JingException e) {
                throw e;
            }
            catch (Throwable t) {
                throw new JingException(t);
            }

            // 重新生成
            BigDecimal maxDuration = Const.NUM_0;
            BigDecimal duration = Const.NUM_0;
            BigDecimal temp;
            for (M3U8MediaPartition p$ : m3u8.getMediaPartitions()) {
                duration = duration.add(temp = BigDecimal.valueOf(p$.getDuration()));
                if (maxDuration.compareTo(temp) < 0) {
                    maxDuration = temp;
                }
            }
            m3u8.getMediaPartitions().clear();
            M3U8MediaPartition partition = new M3U8MediaPartition();
            partition.setDuration(duration.doubleValue());
            partition.setMedia(tsFile.getName());
            m3u8.getMediaPartitions().add(partition);
            m3u8.getTargetDuration().setTargetDuration(maxDuration.doubleValue());
        }
    }

    private static byte[] request(M3U8Config config, String url) {
        LOGGER.debug("request for m3u8 [{}]", url);
        Sender sender = config.getSender();
        int retry = config.getRetry();
        int cnt = 1;
        Wrapper<String> stringWrapper = new Wrapper<>();
        do {
            try {
                WebHeaders headers = new WebHeaders();
                headers.set(WebConst.HEADER_USER_AGENT, WebConst.HEADER_USER_AGENT_VALUE);
                RequestEntity<?> requestEntity = new RequestEntity<>(headers);
                requestEntity.setProxy(config.getProxy());
                ResponseEntity<byte[]> responseEntity = sender.get(url, requestEntity, byte[].class);
                if (!responseEntity.isSuccess()) {
                    throw new JingException(String.valueOf(responseEntity.getStatusCode()));
                }
                return responseEntity.getBody();
            }
            catch (Throwable t) {
                LOGGER.warn(t, "request failed {}/{}", cnt, retry);
                cnt ++;
            }
        }
        while (cnt <= retry);
        throw new JingException("failed to request [{}]", url);
    }

    private static String mergeUrl(String start, String end) {
        if (end.startsWith("/")) {
            end = end.replaceFirst("/", "");
        }
        int position = 0;
        String subEnd, tempEnd = end;
        while ((position = end.indexOf("/", position)) != -1) {
            subEnd = end.substring(0, position + 1);
            if (start.endsWith(subEnd)) {
                tempEnd = end.replaceFirst(subEnd, "");
                break;
            }
            ++position;
        }
        return start + tempEnd;
    }

    public static void outputM3U8(M3U8 m3u8) {
        try (FileContent fileContent = new FileContent()) {
            // 头
            fileContent.add('#' + M3U8TagType.EXT_BEGIN.getTagName());
            if (null != m3u8.getVersion()) fileContent.add('\n' + m3u8.getVersion().output());
            if (null != m3u8.getAllowCache()) fileContent.add('\n' + m3u8.getAllowCache().output());
            if (null != m3u8.getTargetDuration()) fileContent.add('\n' + m3u8.getTargetDuration().output());
            if (null != m3u8.getMediaSequence()) fileContent.add('\n' + m3u8.getMediaSequence().output());
            if (null != m3u8.getPlaylistType()) fileContent.add('\n' + m3u8.getPlaylistType().output());
            if (null != m3u8.getKey()) fileContent.add('\n' + m3u8.getKey().output());
            for (M3U8MediaPartition p$ : m3u8.getMediaPartitions()) {
                fileContent.add('\n' + p$.output());
            }
            fileContent.add("\n#" + M3U8TagType.EXT_END.getTagName());
            fileContent.write(m3u8.getOutputFile());
        }
    }

    private static class AsyncDownloader extends Thread {
        private static final JingLogger LOGGER = JingLogger.getLogger(AsyncDownloader.class);

        private final M3U8Config config;

        private final M3U8 m3u8;

        private AsyncDownloader(M3U8Config config, M3U8 m3u8) {
            this.config = config;
            this.m3u8 = m3u8;
        }

        @Override public void run() {
            LOGGER.info("[{}] m3u8启动异步模式...", m3u8.getOutputFile().getName());
            execute(config, m3u8);
        }
    }

    private static class M3U8MediaPartitionDownloader implements Callable<Void> {
        private final M3U8Config config;

        private final M3U8 m3u8;

        private final List<M3U8MediaPartition> partitions;

        private final File[] mediaFiles;

        private final int mediaIndex;

        private M3U8MediaPartitionDownloader(M3U8Config config, M3U8 m3u8, List<M3U8MediaPartition> partitions, File[] mediaFiles, int mediaIndex) {
            this.config = config;
            this.m3u8 = m3u8;
            this.partitions = partitions;
            this.mediaFiles = mediaFiles;
            this.mediaIndex = mediaIndex;
        }

        @Override
        public Void call() {
            M3U8MediaPartition partition = partitions.get(mediaIndex);
            LOGGER.info("downloading media {}/{} ...", (mediaIndex + 1), partitions.size());
            LOGGER.debug("media: [{}]", partition);
            if (null != config.getMediaPartitionStartListener()) {
                config.getMediaPartitionStartListener().execute(config, m3u8, mediaIndex);
            }
            int retryCount = 0;
            int retryLimit = config.getRetry();
            boolean success = false;

            while (retryCount < retryLimit && !success) {
                try {
                    byte[] buffer = request(config, partition.getMedia());
                    byte[] key = m3u8.getKey().getKey();
                    byte[] iv = m3u8.getKey().getIv();
                    if (null != key && null != iv) {
                        buffer = AES.decrypt(buffer, key, iv);
                    }
                    File tsFile = new File(config.getOutputPath() + File.separator + config.getFileName() + "." + mediaIndex + ".ts");
                    try (FileOutputStream fos = new FileOutputStream(tsFile)) {
                        fos.write(buffer);
                        mediaFiles[mediaIndex] = tsFile;
                        // 覆盖尸块
                        partition.setMedia(tsFile.getName());
                    }
                    catch (Throwable t) {
                        throw new JingException(t);
                    }
                    success = true;
                }
                catch (Exception e) {
                    retryCount ++;
                    LOGGER.error("failed to download , retry condition: {}/{}", retryCount, retryLimit);
                    if (retryCount >= retryLimit) {
                        if (null != config.getMediaPartitionFailListener()) {
                            config.getMediaPartitionFailListener().execute(config, m3u8, mediaIndex);
                        }
                        throw new JingException(e, "failed to download media: {}/{}", (mediaIndex + 1), partitions.size());
                    }
                }
            }
            LOGGER.info("success to download media: {}/{}", (mediaIndex + 1), partitions.size());
            if (null != config.getMediaPartitionSuccessListener()) {
                config.getMediaPartitionSuccessListener().execute(config, m3u8, mediaIndex);
            }
            return null;
        }
    }

    public static void main(String[] args) {
        M3U8Config config = new M3U8Config();
        config.setOutputPath("temp");
        config.setFileName("JUY-809");
        config.setSkipUnknownTags(false);
        config.setAutoCombine(false);
        config.setProxy(new Proxy(Proxy.Type.HTTP, new InetSocketAddress("127.0.0.1", 7890)));
        String m3u8Url = "https://api.6b18726.com/pwa/v2_0/getM3u8.m3u8?v=439137&params=5649716c6b674d6346686f3462316c5573735a5834364f64534457694a336e68356976423445774b59523431516a4e6f7141625073516b62702b704679465a424b4c2f4f796e5476696b4f4542596d755072342b44773d3d";
        M3U8 m3u8 = downloadM3U8(config, m3u8Url);
        LOGGER.debug(m3u8);
    }
}
